'''
Created on Nov 23, 2010

@author: oabalbin
'''
import os
import subprocess
import pysam


def count_covariates(reference_genome, dbsnp_ref, bam_file_name, use_mem, num_cores, path_to_gatk):
    '''
    Generated the file of count covariates, which is required for base quality recalibration
    
    java -Xmx16g -jar /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/GenomeAnalysisTK.jar 
    -l INFO -nt 6 -R /exds/projects/alignment_indexes/gatk/hg19/hg19.fa 
    -B:mask,VCF /exds/projects/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed5.processed.vcf 
    -I s_3_12_sequence.hg19.aln2.rmdup.sorted.bam -T CountCovariates -cov ReadGroupCovariate 
    -cov QualityScoreCovariate -cov CycleCovariate -cov DinucCovariate 
    -recalFile s_3_12_sequence.hg19.aln2.rmdup.sorted.recal_data.csv
    
    
    Note: you can use the --process_nth_locus/ -pN 3 to speed up the process. 
    Which force the program to evaluate every 3 loci, which makes to skip in general ~66 % data
    when dbsnp is also included. Important for Large BAM files
    '''
    
    gatk_command=path_to_gatk+'GenomeAnalysisTK.jar'
    outfile=bam_file_name.replace('.bam','.recal_data.csv')
    args = ['java','-Xmx'+str(use_mem)+'g', '-jar',gatk_command,
            '-l', 'INFO', '-nt', str(num_cores), '-R',reference_genome, 
            '-B:dbsnp,VCF', dbsnp_ref, '-I', bam_file_name, '-T','CountCovariates',
            '-cov', 'ReadGroupCovariate','-cov', 'QualityScoreCovariate', '-cov', 
            'CycleCovariate', '-cov', 'DinucCovariate', '-recalFile',outfile]
    
    retcode = subprocess.call(args)
    
    return outfile

def table_recalibration(reference_genome,bam_file_name,recal_file_name, use_mem, path_to_gatk):
    '''
    Generates a new bam file with recalibrated quality scores
    
    java -Xmx16g -jar /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/GenomeAnalysisTK.jar 
    -l INFO -R /exds/projects/alignment_indexes/gatk/hg19/hg19.fa -I s_3_12_sequence.hg19.aln2.rmdup.sorted.bam 
    -T TableRecalibration -o s_3_12_sequence.aln2.rmdup.sorted.recal.bam 
    -recalFile s_3_12_sequence.hg19.aln2.rmdup.sorted.before.recal_data.csv
    The original function in gatk does not support multiple threads yet
    '''
    gatk_command=path_to_gatk+'GenomeAnalysisTK.jar'
    outfile=bam_file_name.replace('.bam','.recal.bam')
    
    args=['java','-Xmx'+str(use_mem)+'g', '-jar',gatk_command,
            '-l', 'INFO','-R',reference_genome, '-I', bam_file_name, '-T','TableRecalibration',
            '-o',outfile, '-recalFile',recal_file_name]
    
    retcode = subprocess.call(args)
    
    return outfile


def analyze_covariates(recal_file_name,outputdir, resources_folder, rscipt_path, use_mem, path_to_gatk):
    '''
    In this script is important to specify the path to resources in which you have the Rscripts for making the plots and the -R scripts
    -resources /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources (resources_folder)
    -Rscript /exds/sw/local/R-project/bin/Rscript (rscipt_path)
    -ignoreQ 5
    
    java -Xmx16g -jar /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/AnalyzeCovariates.jar 
    -recalFile s_3_12_sequence.hg19.aln2.rmdup.sorted.recal_data.csv 
    -outputDir /exds/users/oabalbin/projects/snps/exomes/aM18/recal_analysis/before/ 
    -resources /exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources/ 
    -Rscript /exds/sw/local/R-project/bin/Rscript
    '''
    
    gatk_command=path_to_gatk+'AnalyzeCovariates.jar'
    args=['java','-Xmx'+str(use_mem)+'g', '-jar',gatk_command,
            '-recalFile',recal_file_name,'-outputDir',outputdir,'-resources',resources_folder,'-Rscript',rscipt_path]
    retcode = subprocess.call(args)
    
    return retcode


def check_create_dir(full_path_name):
        if not os.path.isdir(full_path_name):
            os.mkdir( full_path_name )


def main_baseQ_recalibration(ref_genome, ref_dbsnp, bam_file_name, use_mem, num_cores, 
                             gatk_run_dict, recal_analysis_outputdir):
    '''
    This is a wrap for the base quality recalibration using  gatk (Broad).
    Input: reference genome, vcf dbsnp database, indexed original bam file, 
        memory available, num of cores to use, analysis output directory
        path to gatk, resources folder, and RScript
    
    Returns: recalibrate bam file path, Writes recalibrated bam file and 
        Count covariate tables
        
    gatk quality recalibration steps
    CountCovariates, TableRecalibrate, 
    samtools index on the recalibrated bam file, CountCovariates
    analyze covariates before and after recalibration
    Most of the parameters in this function should go into a dictionary structure.
    '''
    
    path_to_gatk, resources_folder,rscipt_path =\
    gatk_run_dict['path_to_gatk'], gatk_run_dict['resources_folder'], gatk_run_dict['rscipt_path']
    # Generate count covariates table for recalibration
    pre_recal_table_file = count_covariates(ref_genome, ref_dbsnp, bam_file_name, 
                                            use_mem, num_cores, path_to_gatk)
    # Recalibrate original bam file
    #pre_recal_table_file='/exds/users/oabalbin/projects/snps/exomes/aM18/test/s_3_12_sequence.hg19.aln2.rmdup.sorted.recal_data.csv'
    recal_bam_file_name = table_recalibration(ref_genome, bam_file_name,
                                              pre_recal_table_file, use_mem, path_to_gatk)  
    # Index recalibrated bam file
    recal_bam_indexed_name = recal_bam_file_name.replace('.bam','.bam.bai')
    pysam.index(recal_bam_file_name, recal_bam_indexed_name)
    # Generate count covariates table for recalibrated bam file
    #recal_bam_indexed_name='/exds/users/oabalbin/projects/snps/exomes/aM18/test/s_3_12_sequence.hg19.aln2.rmdup.sorted.recal.bam'
    post_recal_table_file = count_covariates(ref_genome, ref_dbsnp, recal_bam_indexed_name, 
                                             use_mem, num_cores, path_to_gatk)
    
    # Analyze covariates before and after recalibration
    before_recal_results_path = recal_analysis_outputdir+'before/'
    after_recal_results_path = recal_analysis_outputdir+'after/'
    
    check_create_dir(before_recal_results_path)
    check_create_dir(after_recal_results_path)
    
    # use the retcode to print to a logfile
    retcode = analyze_covariates(pre_recal_table_file, before_recal_results_path, 
                                 resources_folder, rscipt_path, use_mem, path_to_gatk)
    retcode2 = analyze_covariates(post_recal_table_file, after_recal_results_path, 
                                  resources_folder, rscipt_path, use_mem, path_to_gatk)
    
    return recal_bam_file_name

    
## Test Script
if __name__ == '__main__':
    
    ref_genome='/exds/projects/alignment_indexes/gatk/hg19/hg19.fa'
    ref_dbsnp='/exds/projects/alignment_indexes/gatk/hg19/dbsnp132_00-All_processed5.processed.vcf'
    bam_file_name='/exds/users/oabalbin/projects/snps/exomes/aM18/test/s_3_12_sequence.hg19.aln2.rmdup.sorted.bam'
    use_mem=24
    num_cores=6
    recal_analysis_outputdir='/exds/users/oabalbin/projects/snps/exomes/aM18/test/recal_analysis/'
     
    gatk_run_dict = {'path_to_gatk':'/exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/', 
                     'resources_folder':'/exds/sw/bioinfo/gatk/GenomeAnalysisTK-1.0.4705/resources/', 
                     'rscipt_path':'/exds/sw/local/R-project/bin/Rscript'}
    
        
    main_baseQ_recalibration(ref_genome, ref_dbsnp, bam_file_name, use_mem, num_cores, 
                                 gatk_run_dict, recal_analysis_outputdir)

    
    




    
    
    
    
